use super::types::*;

#[cfg(feature = "linalg")]
use super::util::num_utils;

use num_traits::identities::{One, Zero};
use std::collections::HashMap;

// Not regex::bytes::Regex: used for reading YAML which is always unicode
use regex::Regex;

// parse fractions
use num_bigint::BigUint;
use num_rational::Ratio;
use num_traits::ToPrimitive;
use std::str::FromStr;

/********************************************
 * Grammarfile parsing & grammar validation *
 ********************************************/

pub trait GetOk<'a, I: std::fmt::Display, R> {
    fn get_ok(&self, index: I) -> std::result::Result<&R, anyhow::Error>;
    fn has(&self, index: I) -> bool;
}

/// With useful error message
impl<'a, I> GetOk<'a, I, serde_yaml::Value> for serde_yaml::Value
where
    I: serde_yaml::Index + std::fmt::Display + Clone,
{
    fn get_ok(&self, index: I) -> std::result::Result<&serde_yaml::Value, anyhow::Error> {
        match self.get(&index) {
            Some(value) => Ok(value),
            None => Err(anyhow::Error::msg(format!("no {} in {:?}", index, self))),
        }
    }

    fn has(&self, k: I) -> bool {
        matches!(self.get(k), Some(_))
    }
}

#[derive(Debug, PartialEq, Eq)]
pub enum ValidationResult {
    MissingStartSymbol,
    MissingNonterminal(N, N),
    ProbabilitiesInconsistent { n1: N, is: Prob, should: Prob },
    Ok,
}

use std::fmt;
impl fmt::Display for ValidationResult {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ValidationResult::MissingStartSymbol => {
                write!(f, "Grammar has no rules for nonterminal 0")
            }
            ValidationResult::MissingNonterminal(n1, n2) => {
                write!(f, "Undefined nonterminal {} in a rule of {}", n2, n1)
            }
            ValidationResult::ProbabilitiesInconsistent { n1, is, should } => write!(
                f,
                "Inconsistent probabilities for {}: is {}, should be {}",
                n1, is, should
            ),
            ValidationResult::Ok => write!(f, "Validation OK"),
        }
    }
}

#[derive(Debug, PartialEq)]
pub enum ValidationFindings {
    Unproductive {
        culprits: Vec<N>,
    },
    #[cfg(feature = "linalg")]
    Runaway {
        lambda1: f64,
    },
}

use std::collections::HashSet;

pub fn nonterminals(grammar: &SCFG) -> HashSet<&N> {
    grammar
        .rules
        .iter()
        .map(|(lhs_n, _)| lhs_n)
        .collect::<HashSet<&N>>()
}

/// Check whether grammar has loops that can never resolve to terminal sentences
/// @precondition no epsilon-rules
pub fn detect_looping_nonterminals(grammar: &SCFG) -> HashSet<&N> {
    let all = nonterminals(grammar);

    let mut good: HashSet<&N> = HashSet::new();

    'graph_loop: loop {
        let mut added = HashSet::new();

        'examining: for (lhs_n, rhs_list) in grammar.rules.iter() {
            if good.contains(lhs_n) {
                continue 'examining;
            };

            let rules_good = rhs_list
                .iter()
                .map(|rhs_entry| {
                    // false == There are still non-(known-good) ones in the RHS
                    !rhs_entry
                        .rhs
                        .iter()
                        .any(|x| match x {
                            Symbol::Nonterminal(y) => !good.contains(y),
                            _ => false,
                        })
                })
                .collect::<Vec<bool>>();

            // There exists a rule leading only to terminals or already known good nonterminals, therefore this will also be good (i.e. possible to resolve to terminals).
            if rules_good.iter().any(|x| *x) {
                added.insert(lhs_n);
            }
        }

        if added.is_empty() {
            break 'graph_loop;
        } else {
            good = good.union(&added).cloned().collect();
        }
    }

    all.difference(&good).cloned().collect::<HashSet<&N>>()
}

#[cfg(feature = "linalg")]
pub mod linalg {
    use super::*;
    extern crate nalgebra as na;

    use na::sparse::CsMatrix;
    use na::DVector;
    use nalgebra::base::dimension::Dynamic;

    #[cfg(feature = "linalg")]
    pub fn compute_expansion_matrix(grammar: &SCFG) -> nalgebra::CsMatrix<f64, Dynamic, Dynamic> {
        use itertools::Itertools;

        let nts = nonterminals(grammar);

        // Close gaps in nonterminal numbering, which are allowed to exist
        let mut mapping = HashMap::new();
        for (i, y) in nts.iter().sorted().enumerate() {
            mapping.insert(*y, i);
        }

        let nrows = nts.len();
        let ncols = nts.len();
        let mut triples = vec![];

        for (lhs_n, rhs_list) in grammar.rules.iter().sorted_by(|a, b| Ord::cmp(&a.0, &b.0)) {
            let counts = rhs_list
                .iter()
                .map(|entry| {
                    entry
                        .rhs
                        .iter()
                        .filter_map(|x| match x {
                            Symbol::Nonterminal(y) => Some((
                                BinaryFraction::into_bigrational(entry.probability.clone())
                                    .to_f64()
                                    .unwrap(),
                                *y,
                            )),
                            _ => None,
                        })
                        .collect::<Vec<(f64, N)>>()
                })
                .flatten()
                .fold(HashMap::new(), |mut counts, (r, y)| {
                    *counts.entry(y).or_insert(0.0f64) += r;
                    counts
                });

            for (y, n) in counts {
                triples.push((mapping.get(&y).unwrap(), *mapping.get(lhs_n).unwrap(), n));
            }
        }

        let triples = triples;

        let irows = triples
            .iter()
            .map(|tuple| *tuple.0 as usize)
            .collect::<Vec<_>>();
        let icols = triples
            .iter()
            .map(|tuple| tuple.1 as usize)
            .collect::<Vec<_>>();
        let vals = triples
            .iter()
            .map(|tuple| tuple.2 as f64)
            .collect::<Vec<_>>();

        CsMatrix::from_triplet(nrows, ncols, &irows[..], &icols[..], &vals[..])
    }

    /// @brief Detect runaway behaviour
    /// @return >=1.0 (-error) signals that with nonzero probability, sentential forms will
    /// tend to expand with zero chance of termination, when derived at random.
    ///
    /// Estimate Dominant Eigenvalue of the Nonterminal Population Transition Matrix
    /// It's our way to detect runaway behaviour. If this exceeds 1.0, the grammar's random
    /// derivations are inclined to produce ever longer strings of nonterminals without ever
    /// terminating. That's not a good grammar to be using.
    ///
    /// NB the matrix may very well fail to be mixing / ergodic, for example [[0 2] [0.5 0]].
    /// Periodicity should not be a problem if we first multiply the matrix by itself the
    /// least amount of times that will eliminate all periodicity, lcm(n, ..., 1). Right?
    ///
    pub fn estimate_dominant_eigenvalue(grammar: &SCFG) -> f64 {
        let a: na::Matrix<_, _, _, _> = compute_expansion_matrix(&grammar).into();

        let numbers: Vec<usize> = (1..=a.nrows()).collect();
        let p = num_utils::lcm(&numbers[..]);

        let ap = num_utils::exp_associative(a.clone(), p);

        use core::iter::repeat;
        let mut b: nalgebra::DVector<f64> = DVector::from_iterator(
            a.nrows(),
            vec![1.0f64]
                .into_iter()
                .chain(repeat(0.0f64).take(a.nrows() - 1)),
        );

        let mut i = 0;
        let mut last_norm = 0.0;

        // Power iteration, arbitrarily capped or until it appears to converge
        loop {
            b = ap.clone() * b;
            let norm = b.norm();
            b /= norm;
            if i == 32 || (norm - last_norm).abs() <= 1e-4 {
                break norm;
            }
            last_norm = norm;
            i += 1;
        }
    }
}

// Detect more types of issues (grammar expands too fast; is not deterministic,
// is ambiguous (important but undecidable, can only over-approximate that ...
// no attempt is made at this time - TODO) or contains an unproductive loop)
pub fn further_validate_grammar(grammar: &SCFG) -> Vec<ValidationFindings> {
    let mut findings = vec![];

    let loopies = detect_looping_nonterminals(grammar);
    if !loopies.is_empty() {
        use itertools::Itertools;
        findings.push(ValidationFindings::Unproductive {
            culprits: loopies.iter().copied().sorted().copied().collect(),
        });
    }

    #[cfg(feature = "linalg")]
    {
        let lambda1 = linalg::estimate_dominant_eigenvalue(&grammar);
        // Allow for a certain amount of error in the numerical computation.
        // Future work: determine error tolerance parameters scientifically
        if (lambda1 - 1.0) > 1e-2 {
            findings.push(ValidationFindings::Runaway { lambda1 });
        }
    }

    findings
}

/// Detect non-wellformed grammars.
pub fn syntactically_validate_grammar(grammar: &SCFG) -> ValidationResult {
    if !nonterminals(grammar).iter().any(|x| **x == 0) {
        return ValidationResult::MissingStartSymbol;
    }

    for (lhs_n, rhs_list) in grammar.rules.iter() {
        let mut cumul_lower: Prob = Zero::zero();

        for rhs_entry in rhs_list.iter() {
            #[allow(clippy::float_cmp)]
            if cumul_lower != rhs_entry.cumul_lower {
                return ValidationResult::ProbabilitiesInconsistent {
                    n1: *lhs_n,
                    is: cumul_lower,
                    should: rhs_entry.cumul_lower.clone(),
                };
            }
            cumul_lower += rhs_entry.probability.clone();
            // the "probability" of the highest rule should complete it to 1.0 exactly
            // -> should probably just emit a warning and tweak the rule weights

            for symbol in rhs_entry.rhs.iter() {
                if let Symbol::Nonterminal(n) = symbol {
                    if !grammar.rules.contains_key(n) {
                        return ValidationResult::MissingNonterminal(*lhs_n, *n);
                    }
                }
            }
        }

        #[allow(clippy::float_cmp)]
        if cumul_lower != One::one() {
            return ValidationResult::ProbabilitiesInconsistent {
                n1: *lhs_n,
                is: cumul_lower,
                should: One::one(),
            };
        }
    }

    ValidationResult::Ok
}

pub fn parse_grammar(contents: &str) -> Result<SCFG, anyhow::Error> {
    let res: serde_yaml::Value = serde_yaml::from_str(contents)?;

    let version_seq = res.get_ok("version").ok().and_then(|x| x.as_str());

    match version_seq {
        None => parse_grammar_simple(res),
        Some(version) => match version {
            "0.0.1" => parse_grammar_v1(res),
            _ => Err(anyhow::format_err!("unknown version")),
        },
    }
}

lazy_static! {
    static ref PART: Regex = Regex::new(
        r#"(?x)
           (?:
              "(?P<quoted>(?:[^"\\]|\\.)*)")         # A quoted piece of string
            | (?P<bytestring><[^>]*>)                # Bytes, like <0A> <3F>
            | (?P<fraction>[1-9][0-9]*/[1-9][0-9]*)  # 1/4, 1/3, 255/256 ...
            | (?P<naked>[0-9]+[^\\\ \n"]*)           # 0.25
        "#
    )
    .unwrap();
    static ref HEX: Regex = Regex::new(r#"[0-9A-Fa-f]+"#).unwrap();
}

/// Parse part of a right hand side into a "Symbol" (Nonterminal / Terminals)
fn bean_parse(bean: &regex::Captures) -> Result<Symbol, anyhow::Error> {
    if let Some(naked) = bean.name("naked") {
        match naked.as_str().parse::<u32>() {
            Ok(n) => Ok(Symbol::Nonterminal(n)),
            Err(e) => Err(anyhow::Error::from(e)),
        }
    } else if let Some(bytestring) = bean.name("bytestring") {
        let beans = HEX.captures_iter(bytestring.as_str());
        let xs: Vec<u8> = beans
            .map(|x| u8::from_str_radix(x.get(0).unwrap().as_str(), 16).unwrap())
            .collect();
        Ok(Symbol::Terminals(xs))
    } else if let Some(quoted) = bean.name("quoted") {
        Ok(Symbol::Terminals(quoted.as_str().bytes().collect()))
    } else {
        unreachable!()
    }
}

fn uniform_choices(
    num_rules: usize,
    v: &mut dyn Iterator<Item = Vec<Symbol>>,
) -> Result<Vec<RuleEntry>, anyhow::Error> {
    let mut vec = vec![];
    let mut cumul_lower: Prob = Zero::zero();

    let p_base_approx =
        BinaryFraction::from_float_between_0_and_1(1.0f32 / num_rules as f32).unwrap();

    use core::iter::repeat;
    let e: BinaryFraction = One::one();

    let probabilities = repeat(p_base_approx.clone())
        .take(num_rules - 1)
        .chain(vec![e - p_base_approx * (num_rules - 1)].into_iter());

    for (rhs, rule_prob) in v.zip(probabilities) {
        vec.push(RuleEntry {
            probability: rule_prob.clone(),
            cumul_lower: cumul_lower.clone(),
            rhs,
        });

        cumul_lower += rule_prob;
    }

    Ok(vec)
}

pub fn parse_fraction(fraction: &str) -> Result<BinaryFraction, anyhow::Error> {
    match fraction.split('/').collect::<Vec<_>>()[..] {
        [numer, denom] => {
            let numer = BigUint::from_str(numer)?;
            let denom = BigUint::from_str(denom)?;
            let g = denom.trailing_zeros().unwrap() as i64;
            let denom_g = denom.clone() >> g;
            if denom_g.is_one() {
                // denom is power of 2:
                Ok(BinaryFraction::new(numer, g))
            } else {
                // denom isn't power of 2, approximate as f32:
                let f32_approx = Ratio::new(numer, denom).to_f32().unwrap();
                let ratio = Ratio::from_float(f32_approx).unwrap();
                Ok(BinaryFraction::from_bigrational(Ratio::new_raw(
                    ratio.numer().to_biguint().unwrap(),
                    ratio.denom().to_biguint().unwrap(),
                )))
            }
        }
        _ => Err(anyhow::Error::msg(format!(
            "could not parse fraction {:?}",
            fraction
        ))),
    }
}

pub fn parse_grammar_v1(res: serde_yaml::Value) -> Result<SCFG, anyhow::Error> {
    let rules_seq = res
        .get_ok("rules")?
        .as_mapping()
        .ok_or_else(|| anyhow::Error::msg("rules is not a map"));

    let mut hm = HashMap::new();

    if let Ok(rules_block) = rules_seq {
        for mapping in rules_block.iter() {
            let n: u32 = mapping
                .0
                .as_u64()
                .ok_or_else(|| anyhow::Error::msg(format!("{:?} is not an integer", mapping.0)))?
                .try_into()?;

            if mapping.1.is_mapping() {
                if let Ok(list) = mapping.1.get_ok("uniform-choice") {
                    // List would be the natural choice, but due to pseudo-intuitive YAML suckage (double quotes
                    // are sometimes part of the string, sometimes not) we have to use > instead.
                    let string = list.as_str();
                    if let Some(string) = string {
                        let beans = PART.captures_iter(string);
                        let mut vec = vec![];
                        for bean in beans {
                            let parsed = bean_parse(&bean)?;
                            // one each - would like to make this more flexible by allowing lists of lists,
                            // but with YAML this is way too arduous and fragile.
                            vec.push(vec![parsed]);
                        }
                        let num_rules = vec.len();
                        let mut iter = vec.into_iter();
                        let vec = uniform_choices(num_rules, &mut iter)?;
                        hm.insert(n as u32, vec);
                    }
                } else if let Ok(list) = mapping.1.get_ok("uniform") {
                    let v = list.as_sequence();

                    match v {
                        None => {
                            return Err(anyhow::format_err!(
                                "{:?}[\"uniform\"] does not map to a list",
                                mapping.0
                            ))
                        }

                        Some(v) => {
                            let mut vec = vec![];

                            for val in v {
                                let mut rhs = vec![];
                                let string = val.as_str().ok_or_else(|| {
                                    anyhow::format_err!("YAML error: value {:?} not a string", val)
                                })?;
                                let beans = PART.captures_iter(string);
                                for bean in beans {
                                    rhs.push(bean_parse(&bean)?);
                                }
                                vec.push(rhs);
                            }

                            let num_rules = vec.len();
                            let mut iter = vec.into_iter();
                            let vec = uniform_choices(num_rules, &mut iter)?;
                            hm.insert(n as u32, vec);
                        }
                    }
                }
            } else {
                let v = mapping.1.as_sequence().ok_or_else(|| {
                    anyhow::Error::msg(format!(
                        "{:?} does not map to a mapping or a list",
                        mapping.0
                    ))
                })?;

                let mut vec = vec![];
                let mut cumul_lower: Prob = Zero::zero();

                for (i, rule_entry) in v.iter().enumerate() {
                    let string = rule_entry.as_str().ok_or_else(|| {
                        anyhow::Error::msg(format!("{:?}[{}] does not map to string", n, i))
                    })?;

                    let mut beans = PART.captures_iter(string);

                    let top_bean = beans
                        .next()
                        .ok_or_else(|| anyhow::Error::msg("empty rule definition"))?;

                    let rule_prob = if let Some(fraction) = top_bean.name("fraction") {
                        parse_fraction(fraction.as_str())?
                    } else if let Some(naked) = top_bean.name("naked") {
                        BinaryFraction::from_bigrational(
                            from_float(naked.as_str().parse::<f32>()?).unwrap(),
                        )
                    } else {
                        return Err(anyhow::Error::msg(format!(
                            "should be a float: {:?}",
                            top_bean
                        )));
                    };

                    let mut rhs = vec![];

                    for bean in beans {
                        rhs.push(bean_parse(&bean)?);
                    }

                    vec.push(RuleEntry {
                        probability: rule_prob.clone(),
                        cumul_lower: cumul_lower.clone(),
                        rhs,
                    });
                    cumul_lower += rule_prob;
                }

                hm.insert(n as u32, vec);
            }
        }
    }

    Ok(SCFG { rules: hm })
}

pub fn parse_grammar_simple(res: serde_yaml::Value) -> Result<SCFG, anyhow::Error> {
    let rules_seq = res
        .get_ok("rules")?
        .as_mapping()
        .ok_or_else(|| anyhow::Error::msg("rules is not a map"));

    let mut hm = HashMap::new();

    if let Ok(rules_block) = rules_seq {
        for mapping in rules_block.iter() {
            let n = mapping
                .0
                .as_u64()
                .ok_or_else(|| anyhow::Error::msg(format!("{:?} is not an integer", mapping.0)))?;
            let v = mapping.1.as_sequence().ok_or_else(|| {
                anyhow::Error::msg(format!("{:?} does not map to a list", mapping.0))
            })?;

            let mut vec = vec![];
            let mut cumul_lower: Prob = Zero::zero();

            for (i, rule_entry) in v.iter().enumerate() {
                let string = rule_entry.as_str().ok_or_else(|| {
                    anyhow::Error::msg(format!("{:?}[{}] does not map to string", n, i))
                })?;

                let mut beans = PART.captures_iter(string);

                let top_bean = beans
                    .next()
                    .ok_or_else(|| anyhow::Error::msg("empty rule definition"))?;

                let rule_prob = if let Some(fraction) = top_bean.name("fraction") {
                    parse_fraction(fraction.as_str())?
                } else if let Some(naked) = top_bean.name("naked") {
                    BinaryFraction::from_bigrational(
                        from_float(naked.as_str().parse::<f32>()?).unwrap(),
                    )
                } else {
                    return Err(anyhow::Error::msg(format!(
                        "should be a float: {:?}",
                        top_bean
                    )));
                };

                let mut rhs = vec![];

                for bean in beans {
                    if let Some(naked) = bean.name("naked") {
                        match naked.as_str().parse::<u32>() {
                            Ok(n) => {
                                rhs.push(Symbol::Nonterminal(n));
                            }
                            Err(e) => {
                                return Err(anyhow::Error::from(e));
                            }
                        }
                    }
                    if let Some(bytestring) = bean.name("bytestring") {
                        let beans = HEX.captures_iter(bytestring.as_str());
                        let xs: Vec<u8> = beans
                            .map(|x| u8::from_str_radix(x.get(0).unwrap().as_str(), 16).unwrap())
                            .collect();
                        rhs.push(Symbol::Terminals(xs));
                    }
                    if let Some(quoted) = bean.name("quoted") {
                        rhs.push(Symbol::Terminals(quoted.as_str().bytes().collect()));
                    }
                }

                vec.push(RuleEntry {
                    probability: rule_prob.clone(),
                    cumul_lower: cumul_lower.clone(),
                    rhs,
                });
                cumul_lower += rule_prob;
            }

            hm.insert(n as u32, vec);
        }
    }

    Ok(SCFG { rules: hm })
}
